package com.oscar.springaidemo.ai.config;

import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.node.QuestionClassifierNode;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.oscar.springaidemo.ai.config.nodeType.ExpaderNode;
import com.oscar.springaidemo.ai.config.nodeType.RecordingNode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.HashMap;
import java.util.List;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;

@Configuration
@Slf4j
public class WorkflowAutoConfiguration {

    KeyStrategyFactory keyStrategyFactory = () -> {
        HashMap<String, KeyStrategy> map = new HashMap<>();
        //用户输入
        map.put("query",new ReplaceStrategy());
        map.put("expander_number",new ReplaceStrategy());
        map.put("expander_content",new ReplaceStrategy());
        return map;
    };

    @Bean
    public QuestionClassifierNode feedbackClassifier(ChatClient chatClient){
        // 评价正负分类节点
        QuestionClassifierNode feedbackClassifier = QuestionClassifierNode.builder()
                .chatClient(chatClient)
                .inputTextKey("input")
                .categories(List.of("positive feedback", "negative feedback"))
                .classificationInstructions(
                        List.of("Try to understand the user's feeling when he/she is giving the feedback."))
                .build();
        return feedbackClassifier;
    }

    @Bean
    public QuestionClassifierNode specificQuestionClassifier(ChatClient chatClient) {
        // 负面评价具体问题分类节点
        QuestionClassifierNode specificQuestionClassifier = QuestionClassifierNode.builder()
                .chatClient(chatClient)
                .inputTextKey("input")
                .categories(List.of("after-sale service", "transportation", "product quality", "others"))
                .classificationInstructions(List.of(
                        "What kind of service or help the customer is trying to get from us? " +
                                "Classify the question based on your understanding."))
                //
                .build();
        return specificQuestionClassifier;
    }


    @Bean(name = "workflowGraph")
    public StateGraph stateGraph(ChatClient defaultChatClient) throws GraphStateException {
        RecordingNode recorderNode = new RecordingNode();
        ExpaderNode expaderNode = new ExpaderNode(defaultChatClient);

        StateGraph stateGraph = new StateGraph("Consumer Service Workflow Demo", keyStrategyFactory)
                .addNode("expander", node_async(expaderNode))
                .addEdge(START,"expander")
                .addEdge("expander", END);

        GraphRepresentation representation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML
                , "expender flow");
        log.info("=====uml flow start=====");
        log.info(representation.content());
        log.info("=====uml flow end=====");


        return stateGraph;
    }


}





